# encoding=utf-8

import os
import sys
import json
import argparse

import threadpool
import tqdm
import threading
from PIL import Image
from PIL import ImageFile

from .. import common
from .. import annotation_utils

ImageFile.LOAD_TRUNCATED_IMAGES = True

_RES = {}


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("image_folder", default="", help="Image folder")
    parser.add_argument("--annotation_folder", default="",
                        help="Annotation folder. This argument will be same "
                             "as image_folder when it's not set.")
    parser.add_argument("output_folder", default="", help="Output folder")
    parser.add_argument("--expand_ratio", type=float, default=0.,
                        help="Expand ratio of bounding box. Default is 0.")
    parser.add_argument("--include_status", default="3,5",
                        help="include status, ',' split if there are more than "
                             "one status. default is '3,5'")
    parser.add_argument("--square", action="store_true", default=False,
                        help="Choose square bounding box to crop.")
    parser.add_argument("--num_thread", type=int, default=16,
                        help="Size of thread pool. Default is 16.")

    args = parser.parse_args()
    if args.annotation_folder == "":
        args.annotation_folder = args.image_folder
    args.include_status = set(args.include_status.split(","))
    if args.num_thread <= 0:
        print("warning: num_thread should be a positive integer!")
        args.num_thread = 16

    return args


FLAG = get_args()


def get_patches(img_path, ann_path):
    name = os.path.basename(img_path)

    image = Image.open(img_path)
    annotation = annotation_utils.read_annotation(ann_path)

    image_width, image_height = image.size

    for obj in annotation.get("object", []):
        ymin, xmin, ymax, xmax = \
            annotation_utils.extract_bndbox_from_object(obj)

        h, w = ymax - ymin, xmax - xmin
        if h <= 0 or w <= 0:
            print("An invalid bounding box in '%s':" % ann_path)
            print(json.dumps(obj["bndbox"], indent=4))
            continue

        if "status" in obj and str(obj["status"]) not in FLAG.include_status:
            continue

        delta_h = int(round(h * FLAG.expand_ratio / 2))
        delta_w = int(round(w * FLAG.expand_ratio / 2))
        ymi, yma = max(0, ymin - delta_h), min(image_height, ymax + delta_h)
        xmi, xma = max(0, xmin - delta_w), min(image_width, xmax + delta_w)

        sub_folder = os.path.join(FLAG.output_folder, obj["name"].strip())
        if not os.path.isdir(sub_folder):
            os.makedirs(sub_folder)
        ttt = "%s_%s_%s_%s.jpg" % (xmin, xmax, ymin, ymax)
        path = os.path.join(sub_folder, name + ttt)
        if FLAG.square:
            cx, cy = (xmi + xma) / 2, (yma + ymi) / 2
            ww, hh = xma - xmi, yma - ymi
            ll_half = max(ww, hh) / 2
            xmi, ymi, xma, yma = cx - ll_half, cy - ll_half, cx + ll_half, cy + ll_half
        p_img = image.crop([xmi, ymi, xma, yma])
        p_img.save(path)

    if _RES["lock"].acquire():
        _RES["process"].update(1)
        _RES["lock"].release()


def create_all_class_folder(root_dir, annotation_paths):
    """ Create all sub-class folder

    :param root_dir: root directory
    :param annotation_paths: a list of annotation path
    :return:
    """
    class_set = set()
    for pt in annotation_paths:
        annotation = annotation_utils.read_annotation(pt)
        for obj in annotation.get("object", []):
            if obj["name"] not in class_set:
                class_set.add(obj["name"])
                sub_folder = os.path.join(root_dir, obj["name"].strip())
                if not os.path.isdir(sub_folder):
                    os.makedirs(sub_folder)
                    # print("Created class folder: ", sub_folder)
                    sys.stdout.write("\r" + " " * 80)
                    sys.stdout.write("\rCreated class folder: {fld}".format(fld=sub_folder))
                    sys.stdout.flush()


def main():
    img_paths = common.find_file_recursive(FLAG.image_folder, ("jpg", "png", "jpeg"),
                                           ignore_case=True)
    ann_paths = common.find_file_recursive(FLAG.annotation_folder, "xml",
                                           ignore_case=True)
    img_paths, ann_paths = common.align_paths(img_paths, ann_paths)

    # create all sub-class folder
    create_all_class_folder(FLAG.output_folder, ann_paths)

    print()
    pool = threadpool.ThreadPool(FLAG.num_thread)
    params = [([x, y], {}) for x, y in zip(img_paths, ann_paths)]
    requests = threadpool.makeRequests(get_patches, params)
    _RES["process"] = tqdm.tqdm(requests)
    _RES["lock"] = threading.Lock()
    # map(pool.putRequest, requests)
    for req in requests:
        pool.putRequest(req)
    pool.wait()
    _RES["process"].refresh()


if __name__ == '__main__':
    main()
